import math
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from matplotlib.patches import Rectangle
from io import BytesIO
import cv2
from PIL import Image
from matplotlib.colors import LinearSegmentedColormap
import matplotlib.patches as patches
def draw_car_position(map_shape, car_size = (666, 472, 294, 513)):

    grid_size =  1 / 0.05
    w,h = map_shape
    y_top = int(h // 2 - 513 / 1000 * grid_size)
    y_bottom = int(h // 2 + (666 - 513) / 1000 * grid_size)
    x_left = int(w // 2 - 472 / 1000 / 2 * grid_size )
    x_right = int(w // 2 + 472 / 1000 / 2 * grid_size )
    # 绘制四条直线表示小车的位置
    plt.plot([x_left, x_right], [y_top, y_top], color='red', linewidth=2)  # 上水平线
    plt.plot([x_left, x_right], [y_bottom, y_bottom], color='red', linewidth=2)  # 下水平线
    plt.plot([x_left, x_left], [y_top, y_bottom], color='red', linewidth=2)  # 左垂直线
    plt.plot([x_right, x_right], [y_top, y_bottom], color='red', linewidth=2)  # 右垂直线


    # 使用 Rectangle 绘制矩形表示小车的位置
    plt.plot(h//2 , w //2, 'ro')  # 'ro' 表示红色的圆点


def draw_rotate_car_position(map_shape, fusepose_rpy , car_size = (666, 472, 294, 513)):
    grid_size =  1 / 0.05
    w,h = map_shape

    # y_bottom = int(513 / 1000 * grid_size)
    # y_top = int((666 - 513) / 1000 * grid_size)
    # x_left = int(- 472 / 1000 / 2 * grid_size )
    # x_right = int(472 / 1000 / 2 * grid_size )

    x_top = int((666 - 513) / 1000 * grid_size) * (-1)
    x_bottom = int(513 / 1000 * grid_size)
    y_left = int(- 472 / 1000 / 2 * grid_size )
    y_right = int(472 / 1000 / 2 * grid_size )

    p1 = np.array([x_top, y_left])
    p2 = np.array([x_top, y_right])
    p3 = np.array([x_bottom, y_right])
    p4 = np.array([x_bottom, y_left])

    # p1 = np.array([x_left_hat, y_top_hat])
    # p2 = np.array([x_right_hat, y_top_hat])
    # p3 = np.array([x_right_hat, y_bottom_hat])
    # p4 = np.array([x_left_hat, y_bottom_hat])

    _, _, yaw = fusepose_rpy
    print(f'yaw {yaw}')

    rotation_matrix = np.array([
            [np.cos(yaw), -np.sin(yaw)],
            [np.sin(yaw), np.cos(yaw)]
        ])

        # 创建局部坐标的列向量

        # 应用旋转矩阵

    delta_size = h // 2
    # x要换成y
    p12, p11  = (np.dot(rotation_matrix, p1) + delta_size).tolist()
    p22, p21  = (np.dot(rotation_matrix, p2) + delta_size).tolist()
    p32, p31  = (np.dot(rotation_matrix, p3) + delta_size).tolist()
    p42, p41  = (np.dot(rotation_matrix, p4) + delta_size).tolist()

    # p11, p12  = (np.dot(rotation_matrix, p1) + delta_size).tolist()
    # p21, p22  = (np.dot(rotation_matrix, p2) + delta_size).tolist()
    # p31, p32  = (np.dot(rotation_matrix, p3) + delta_size).tolist()
    # p41, p42  = (np.dot(rotation_matrix, p4) + delta_size).tolist()




    plt.plot([p11, p21], [p12, p22], color='red', linewidth=2)  # 上水平线
    plt.plot([p21, p31], [p22, p32], color='red', linewidth=2)  # 下水平线
    plt.plot([p31, p41], [p32, p42], color='red', linewidth=2)  # 左垂直线
    plt.plot([p41, p11], [p42, p12], color='red', linewidth=2)  # 右垂直线

    plt.plot(h//2 , w //2, 'ro')  # 'ro' 表示红色的圆点

                             

def visualize_line1(map_grid, fusepose_rpy, index, conf = 0.5):
    # 创建一个自定义颜色映射，从黑色到白色
    # cmap = LinearSegmentedColormap.from_list('custom_cmap', [(0, 'black'), (0.6, 'white'), (1, 'black')])
    # cmap = LinearSegmentedColormap.from_list('custom_cmap', [(0, 'white'), (0.5, 'white'), (1, 'black')])
    cmap = LinearSegmentedColormap.from_list('custom_cmap', [(0, 'white'), (conf, 'white'), (1, 'black')])
    # 使用 plt.imshow 绘制图像
    im = plt.imshow(map_grid, cmap=cmap, vmin=0, vmax=1)  # 自定义映射，vmin和vmax设置为0到1

    # 添加颜色条，标注为 'Occupancy Probability'
    plt.colorbar(im, label='Occupancy Probability')

    # 设置标题
    plt.title(f'Grid Map After Observation {index + 1}')

    # 设置X轴和Y轴标签
    plt.xlabel('X')
    plt.ylabel('Y')
    
    # draw_car_position(map_grid.shape)
    draw_rotate_car_position(map_grid.shape, fusepose_rpy)
    # 保存图像到文件，文件名包含索引
    plt.savefig(f'mid_res/grid_map_observation_{index + 1}.png', bbox_inches='tight')
    plt.close()  # 关闭当前图像，避免重叠



def visualize_line2(map_grid, fusepose_rpy, index):
    # 创建一个自定义颜色映射，从浅灰到白色
    cmap = LinearSegmentedColormap.from_list('custom_cmap', [(0, 'lightgray'), (1, 'white')])

    # 使用 plt.imshow 绘制图像
    im = plt.imshow(map_grid, cmap=cmap, vmin=0, vmax=1)  # 自定义映射，vmin和vmax设置为0到1

    # 添加颜色条，标注为 'Occupancy Probability'
    plt.colorbar(im, label='Occupancy Probability')

    # 设置标题
    plt.title(f'Grid Map After Observation {index + 1}')

    # 设置X轴和Y轴标签
    plt.xlabel('X')
    plt.ylabel('Y')

    # 绘制每个栅格的边框
    grid_size = 1
    for y in range(map_grid.shape[0]):
        for x in range(map_grid.shape[1]):
            # 创建矩形补丁
            rect = patches.Rectangle((x - 0.5, y - 0.5), grid_size, grid_size, linewidth=0.5, edgecolor='gray', facecolor='none')
            # 添加矩形到当前图像
            plt.gca().add_patch(rect)

            # 如果值大于0.5，用浅红色填充
            if map_grid[y, x] > 0.5:
                rect = patches.Rectangle((x - 0.5, y - 0.5), grid_size, grid_size, linewidth=0, facecolor='lightcoral', alpha=0.5)
                plt.gca().add_patch(rect)

    # 绘制旋转后的车位置
    draw_rotate_car_position(map_grid.shape, fusepose_rpy)

    # 保存图像到文件，文件名包含索引
    plt.savefig(f'grid_map_observation_{index + 1}.png', bbox_inches='tight')
    plt.close()  # 关闭当前图像，避免重叠
    plt.close()  # 关闭当前图像，避免重叠



def visualize_video(map_grid, fusepose_rpy, index, conf = 0.5):
    cmap = LinearSegmentedColormap.from_list('custom_cmap', [(0, 'white'), (conf, 'white'), (1, 'black')])

    # 创建一个新的图形对象
    fig, ax = plt.subplots()

    # 使用 plt.imshow 绘制图像
    im = ax.imshow(map_grid, cmap=cmap, vmin=0, vmax=1)  # 自定义映射，vmin和vmax设置为0到1

    # 添加颜色条，标注为 'Occupancy Probability'
    plt.colorbar(im, ax=ax, label='Occupancy Probability')

    # 设置标题
    ax.set_title(f'Grid Map After Observation {index + 1}')

    # 设置X轴和Y轴标签
    ax.set_xlabel('X')
    ax.set_ylabel('Y')

    # 绘制车辆位置
    draw_rotate_car_position(map_grid.shape, fusepose_rpy)

    # 使用 BytesIO 来保存图像
    buf = BytesIO()
    fig.savefig(buf, format='png', bbox_inches='tight')  # 使用 fig.savefig 而不是 plt.savefig
    plt.close(fig)  # 关闭 fig 而不是 plt
    buf.seek(0)  # 将指针移动到开始位置，以便读取图像数据

    # 检查 BytesIO 是否为空
    if not buf.getvalue():
        raise ValueError("BytesIO buffer is empty. Image may not have been saved correctly.")

    # 将 BytesIO 对象转换为 NumPy 数组
    image = Image.open(buf)
    image_np = np.array(image)

    # 如果图像是 RGBA 格式，转换为 BGR 格式以便 OpenCV 使用
    if image_np.shape[2] == 4:
        image_np = cv2.cvtColor(image_np, cv2.COLOR_RGBA2BGR)
    else:
        image_np = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)

    return image_np

def visualize_line(map_grid, index):
    # 创建一个自定义颜色映射，从黑色到白色
    cmap = LinearSegmentedColormap.from_list('custom_cmap', [(0, 'black'), (0.5, 'white'), (1, 'black')])

    # 创建图像
    fig, ax = plt.subplots()
    im = ax.imshow(map_grid, cmap=cmap, vmin=0, vmax=1)  # 自定义映射，vmin和vmax设置为0到1

    # 添加颜色条，标注为 'Occupancy Probability'
    plt.colorbar(im, label='Occupancy Probability')

    # 设置标题
    plt.title(f'Grid Map After Observation {index + 1}')

    # 设置X轴和Y轴标签
    plt.xlabel('X')
    plt.ylabel('Y')
    
    draw_car_position(map_grid.shape)

    # 使用 BytesIO 来保存图像
    buf = BytesIO()
    fig.savefig(buf, format='png', bbox_inches='tight')  # 使用 fig.savefig 而不是 plt.savefig
    plt.close(fig)  # 关闭 fig 而不是 plt
    buf.seek(0)  # 将指针移动到开始位置，以便读取图像数据

    # 检查 BytesIO 是否为空
    if not buf.getvalue():
        raise ValueError("BytesIO buffer is empty. Image may not have been saved correctly.")

    # 将 BytesIO 对象转换为 NumPy 数组
    image = Image.open(buf)
    image_np = np.array(image)

    # 如果图像是 RGBA 格式，转换为 BGR 格式以便 OpenCV 使用
    if image_np.shape[2] == 4:
        image_np = cv2.cvtColor(image_np, cv2.COLOR_RGBA2BGR)
    else:
        image_np = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
    
    return image_np



if __name__ == '__main__':
    map_shape = 100,100
    fusepose_rpy = 0,0, 0
    draw_rotate_car_position(map_shape, fusepose_rpy , car_size = (666, 472, 294, 513))
